import numpy as np
import cupy as cp
import matplotlib.pyplot as plt


# 通用参数
Re = 6371e3  # 地球半径
c = 3e8  # 光速
G = 6.67e-11  # 引力常量
M = 5.965e24  # 地球质量


class Signal:
    def __init__(self, parameter):
        self.Nx = parameter['Nx']
        self.Ny = parameter['Ny']
        self.Wt = parameter['Wt']
        self.Wr = parameter['Wr']
        self.dx = parameter['Dx']
        self.dy = parameter['Dy']
        self.f0 = parameter['F0']
        self.lamda = parameter['Lamda']

        self.H = parameter['H']
        self.alpha_c = parameter['AlphaC']
        self.theta_sq_c = parameter['ThetaSqC']

        self.alpha_c_r = np.deg2rad(self.alpha_c)
        self.theta_sq_c_r = np.deg2rad(self.theta_sq_c)

        self.alpha_scan_min = parameter['alphaScanMin']
        self.alpha_scan_max = parameter['alphaScanMax']
        self.alpha_scan_max_r = np.deg2rad(self.alpha_scan_max)
        self.alpha_scan_min_r = np.deg2rad(self.alpha_scan_min)

        self.PRF = parameter['PRF']
        self.PRT = 1 / self.PRF
        self.Tr = parameter['Tr']
        self.Br = parameter['Br']
        self.Kr = self.Br / self.Tr
        self.Fr = parameter['Fr']

        self.vs = (G * M / (self.H + Re)) ** 0.5
        self.ws = self.vs / (self.H + Re)
        self.theta_c_r = self.cal_theta(self.alpha_c_r, self.theta_sq_c_r)
        self.eta_c_r = self.theta2r_eta(self.theta_c_r)[1]
        self.vg = Re * self.ws * np.cos(self.eta_c_r)
        self.vr = np.sqrt(self.vg * self.vs)

        self.theta_bw_r = 0.886 / self.dy * self.lamda / self.Ny / 4
        self.theta_bw = np.rad2deg(self.theta_bw_r)

        # 主信号区
        self.theta_min_r = self.cal_theta(self.alpha_scan_min_r, self.theta_sq_c_r)
        self.theta_max_r = self.cal_theta(self.alpha_scan_max_r, self.theta_sq_c_r)
        self.theta_min = np.rad2deg(self.theta_min_r)
        self.theta_max = np.rad2deg(self.theta_max_r)

        self.R0_target = self.theta2r_eta(self.alpha_c_r)[0]
        self.R_min = self.R0_target - 5000  # 成像最近斜距
        self.R_max = self.R0_target + 5000  # 成像最远斜距
        self.Rc = self.theta2r_eta(self.theta_c_r)[0]

        self.tr_min = self.R_min * 2 / np.cos(self.theta_sq_c_r) / c
        self.tr_max = self.R_max * 2 / np.cos(self.theta_sq_c_r) / c
        self.tr_w = self.tr_max - self.tr_min

        self.Ta = 0.886 * self.lamda * self.Rc / self.Ny / 4 / self.dy / np.cos(self.theta_sq_c_r) / self.vg
        # self.X = self.Ta * self.vg
        # self.Ta * self.PRF

        # self.t_fast = np.linspace(self.tr_min, self.tr_max + self.Tr, int((self.tr_max + self.Tr - self.tr_min) * self.Fr))
        self.t_fast = np.arange(self.tr_min, self.tr_max + self.Tr, 1/self.Fr)
        self.theta_sq_max_r = self.theta_sq_c_r + self.theta_bw_r / 2
        self.theta_sq_min_r = self.theta_sq_c_r - self.theta_bw_r / 2

        self.theta_radar_max_r = self.cal_theta(self.alpha_c_r, self.theta_sq_max_r)
        self.theta_radar_min_r = self.cal_theta(self.alpha_c_r, self.theta_sq_min_r)
        self.R_radar_min = self.R0_target / np.cos(self.theta_sq_min_r)
        self.R_radar_max = self.R0_target / np.cos(self.theta_sq_max_r)

        self.t_min = -np.sqrt(self.R_radar_min ** 2 - self.R0_target ** 2) / self.vr + 1000 * self.PRT
        self.t_max = -np.sqrt(self.R_radar_max ** 2 - self.R0_target ** 2) / self.vr - 1000 * self.PRT
        # self.nPRF = int(np.ceil((self.t_min - self.t_max) * self.PRF))
        # self.t_slow = np.linspace(self.t_max, self.t_max + self.nPRF * self.PRT, self.nPRF)
        self.t_slow = np.arange(self.t_max, self.t_min, self.PRT)

        # 雷达姿态
        self.A = np.arctan(np.tan(self.alpha_c_r) / np.cos(self.theta_sq_c_r))
        self.theta_yaw = np.arctan(np.sin(self.alpha_c_r) * np.tan(self.theta_sq_c_r))
        self.theta_pit = np.arcsin(np.cos(self.alpha_c_r) * np.sin(self.theta_sq_c_r))

    def signal_main(self):
        tfa, tsw = np.meshgrid(self.t_fast, self.t_slow)
        r_target = np.sqrt(self.R0_target ** 2 + self.vr ** 2 * self.t_slow ** 2)
        # theta_target = self.r2eta_theta(r_target)[1]
        # theta_sq = self.cal_theta_sq(self.alpha_c_r, theta_target)
        theta_sq = np.arccos(self.R0_target / r_target)
        p = self.pattern_double(self.Wt, self.Wr, self.alpha_c_r, theta_sq)
        tfa, r_target = np.meshgrid(self.t_fast, r_target)
        tfa, p = np.meshgrid(self.t_fast, p)
        window = abs(tfa - 2 * r_target / c) < self.Tr/2
        s0 = window * p * np.exp(-1j * 4 * np.pi * self.f0 * r_target / c) * \
             np.exp(1j * np.pi * self.Kr * ((tfa - 2 * r_target) / c) ** 2)
        # pp = window * p
        # pp /= np.max(pp)

        return

    def steer_2d(self, alpha, theta_sq):
        def phase(al, th):
            phx = ((np.cos(self.A) * np.cos(self.theta_yaw) +
                    np.sin(self.A) * np.sin(self.theta_yaw) * np.sin(self.theta_pit)) * np.sin(al) * np.cos(th) +
                   (np.cos(self.A) * np.sin(self.theta_yaw) -
                    np.sin(self.A) * np.cos(self.theta_yaw) * np.sin(self.theta_pit)) * np.sin(th) -
                   np.sin(self.A) * np.cos(self.theta_pit) * np.cos(al) * np.cos(th)) * \
                  self.dx * 2j * np.pi / self.lamda

            phy = (-np.cos(self.theta_pit) * np.sin(self.theta_yaw) * np.sin(al) * np.cos(th) +
                   np.cos(self.theta_pit) * np.cos(self.theta_yaw) * np.sin(th) -
                   np.sin(self.theta_pit) * np.cos(al) * np.cos(th)) * self.dy * 2j * np.pi / self.lamda

            phz = ((np.sin(self.A) * np.cos(self.theta_yaw) -
                    np.cos(self.A) * np.sin(self.theta_yaw) * np.sin(self.theta_pit)) * np.sin(al) * np.cos(th) +
                   (np.sin(self.A) * np.sin(self.theta_yaw) +
                    np.cos(self.A) * np.cos(self.theta_yaw) * np.sin(self.theta_pit)) * np.sin(th) +
                   np.cos(self.A) * np.cos(self.theta_pit) * np.cos(al) * np.cos(th))
            return phx, phy, phz

        phase_x, phase_y, phase_z = phase(alpha, theta_sq)

        shape_x = list(phase_x.shape)
        shape_y = list(phase_y.shape)
        shape_x.extend([self.Nx, 1])
        shape_y.extend([self.Ny, 1])
        sx = np.zeros(tuple(shape_x), dtype='complex')
        sy = np.zeros(tuple(shape_y), dtype='complex')
        p = np.zeros(phase_x.shape, dtype='complex')
        for i in range(4):
            p += np.exp(i * phase_y)
        for i in range(self.Nx):
            sx[..., i, 0] = np.exp(i * phase_x)
        for i in range(self.Ny):
            sy[..., i, 0] = p * np.exp(4 * i * phase_y)

        loc = np.where(phase_z < 0)
        if len(loc) == 2:
            sx[loc[0], loc[1], ...] = np.nan + np.nan * 1j
            sy[loc[0], loc[1], ...] = np.nan + np.nan * 1j
        elif len(loc) == 1:
            sx[loc[0], ...] = np.nan + np.nan * 1j
            sy[loc[0], ...] = np.nan + np.nan * 1j
        return sx, sy

    def vector_2d(self, alpha, theta_sq):
        sx, sy = self.steer_2d(alpha, theta_sq)
        v = self.vector(sx, sy)
        shape = list(theta_sq.shape)
        shape.extend([self.Nx * self.Ny, 1])
        v = v.reshape(shape)
        return v

    def pattern_double(self, wt, wr, alpha, theta_sq):
        v = self.vector_2d(alpha, theta_sq)
        wt = wt.reshape(-1, 1)
        wr = wr.reshape(-1, 1)
        pt = abs(wt.T @ v).reshape(theta_sq.shape)
        pr = abs(wr.T @ v).reshape(theta_sq.shape)
        p = pt * pr
        return p

    def r2eta_theta(self, r):  # 由斜距计算入射角，星下点角
        theta = np.arccos(np.round((r ** 2 + (Re + self.H) ** 2 - Re ** 2) / 2 / (Re + self.H) / r, 12))
        # alpha = np.arccos((r ** 2 + (Re + self.H) ** 2 - Re ** 2) / 2 / (Re + self.H) / r)
        eta = np.arcsin(np.round(np.sin(theta) * (Re + self.H) / Re, 12))
        # eta = np.arcsin(np.sin(alpha) * (Re + self.H) / Re)
        return eta, theta

    def theta2r_eta(self, theta):  # 由星下点角计算斜距、入射角
        eta = np.arcsin(np.round(np.sin(theta) * (Re + self.H) / Re, 12))
        # eta = np.arcsin(np.sin(alpha) * (Re + self.H) / Re)
        beta = eta - theta
        r = np.sin(beta) / np.sin(theta) * Re
        if type(r) == np.float64 and r == 0:
            r = self.H
        elif type(r) == np.array:
            r[np.where(beta == 0)] = self.H
        return r, eta

    @staticmethod
    def cal_alpha_all(theta_sq, theta):
        theta_sq, theta = np.meshgrid(theta_sq, theta)
        return np.arccos(np.cos(theta) / np.cos(theta_sq)), theta_sq, theta

    @staticmethod
    def vector(sx, sy):
        shape = list(range(len(sy.shape)))
        shape.insert(-1, shape.pop())
        sy = np.transpose(sy, shape)
        v = sx @ sy
        return v

    @staticmethod
    # parameter from 《Handbook of Radar Scattering Statistics for Terrain>> Appendix C grass L band HV
    def cal_scatter(eta):  # input: radian output: power
        # HH
        p1 = -29.235
        p2 = 37.550
        p3 = 2.332
        p4 = -2.615
        p5 = 5.0
        p6 = -1.616

        # HV
        # p1 = -40.166
        # p2 = 26.833
        # p3 = 2.029
        # p4 = -1.473
        # p5 = 3.738
        # p6 = -1.324

        sigma0_db = p1 + p2 * np.exp(-p3 * eta) + p4 * np.cos(p5 * eta + p6)
        sigma0 = 10 ** (sigma0_db / 10)
        return sigma0

    @staticmethod
    def cal_theta(alpha, theta_sq):  # input:radian output:radian
        return np.arccos(np.cos(alpha) * np.cos(theta_sq))

    @staticmethod
    def cal_theta_sq(alpha, theta):
        return np.arccos(np.cos(theta) / np.cos(alpha))


if __name__ == '__main__':
    Parameter = {
        'Wt': np.ones((22*16, 1), dtype='complex'),
        'Wr': np.ones((22*16, 1), dtype='complex'),
        'Nx': 22,
        'Ny': 16,
        'Lamda': c / 1.26e9,
        'Dx': 0.65 * c / 1.26e9,
        'Dy': 0.64 * c / 1.26e9,
        'dxRatio': 0.65,
        'dyRatio': 0.64,
        'H': 607e3,
        'F0': 1.26e9,
        'PRF': 3250,  # test
        'AlphaC': 26,
        'ThetaSqC': 15,
        'alphaScanMin': 25.06,
        'alphaScanMax': 27.33,
        'Tr': 20e-6,
        'Br': 60e6,
        'Fr': 90e6
    }
    Sig = Signal(Parameter)
    Sig.signal_main()
    print('End')

